package _13_red_black_tree.case4;

import static _13_red_black_tree.case4.RedBlackTree.Color.BLACK;
import static _13_red_black_tree.case4.RedBlackTree.Color.RED;

/**
 * @author: mornd
 * @dateTime: 2023/6/17 - 22:24
 * 红黑树定义
 * 第三次
 */
class RedBlackTree {
    Node root;

    enum Color {
        RED, BLACK;
    }

    /**
     * 红黑树节点
     */
    static class Node {
        int key;
        Object value;
        Node left;
        Node right;
        Node parent;
        Color color = RED; // 默认红色

        public Node(int key) {
            this.key = key;
        }

        public Node(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public Node(int key, Color color) {
            this.key = key;
            this.color = color;
        }

        public Node(int key, Object value, Color color) {
            this.key = key;
            this.value = value;
            this.color = color;
        }

        public Node(int key, Color color, Node left, Node right) {
            this.key = key;
            this.color = color;
            this.left = left;
            this.right = right;
            if(left != null) {
                left.parent = this;
            }
            if(right != null) {
                right.parent = this;
            }
        }

        // 判断当前节点是否是左孩子
        private boolean isLeftChild() {
            return parent != null && parent.left == this;
        }

        // 获取叔叔节点（父节点的兄弟节点）
        private Node uncle() {
            if (parent == null || parent.parent == null) {
                return null;
            }
            if (parent.isLeftChild()) {
                return parent.parent.right;
            } else {
                return parent.parent.left;
            }
        }

        // 获取兄弟节点（同一个父亲）
        private Node sibling() {
            if (parent == null) {
                return null;
            }
            return isLeftChild() ? parent.right : parent.left;
        }
    }

    // 删除
    public void remove(int k) {
        Node deleted = find(k);
        if(deleted != null) {
            remove(deleted);
        }
    }

    /**
     * 删除
     * case0：删除节点有两个孩子 => 复制后继节点的k、v，然后简化成只有一个或没有孩子
     * case1：删除的是根节点，又分为 根有孩子 => 孩子替死    根没有孩子 => root=null
     * 如果删除的是红色，直接删除，不需要调整
     * case2：删的是黑，剩下的是红 => 剩下的变黑即可
     * 此时已经触发双黑
     * case3：被调整节点的兄弟为红 =>  父亲左\右旋转一次，父亲变红，兄弟变黑，再转化为case4或case5
     * case4：被调整节点的兄弟为黑，兄弟的两个孩子都是黑
     *        => 先将兄弟变红，此时如果父亲是红，则将父亲变黑即可，反之让父亲再次触发双黑
     * case5：被调整节点的兄弟为黑，兄弟的两个孩子有一个或都是红
     *         => 此时又分4种情况
     *              兄弟是左孩子，左侄子是红色 LL
     *              兄弟是左孩子，右侄子是红色 LR
     *              兄弟是右孩子，左侄子是红色 RL
     *              兄弟是右孩子，右侄子是红色 RR
     */
    private void remove(Node deleted) {
        Node replaced = findReplaced(deleted);
        Node parent = deleted.parent;

        if (replaced == null) {
            if (root == deleted) {
                root = null;
            } else {
                if (isBlack(deleted)) {
                    fixDoubleBlack(deleted);
                }

                if (deleted.isLeftChild()) {
                    parent.left = null;
                } else {
                    parent.right = null;
                }
                deleted.parent = null;
            }
            return;
        }

        if (deleted.left == null || deleted.right == null) {
            if (root == deleted) {
                root.key = replaced.key;
                root.value = replaced.value;
                root.left = root.right = null;
            } else {
                if (deleted.isLeftChild()) {
                    parent.left = replaced;
                } else {
                    parent.right = replaced;
                }
                replaced.parent = parent;
                deleted.left = deleted.right = deleted.parent = null;

                if (isBlack(deleted)) {
                    if (isBlack(replaced)) {
                        fixDoubleBlack(replaced);
                    } else {
                        replaced.color = BLACK;
                    }
                }
            }
            return;
        }

        deleted.key = replaced.key;
        deleted.value = replaced.value;

        remove(replaced); // 替死

    }

    private void fixDoubleBlack(Node n) { // n必须为黑
        if(n == root) {
            return;
        }
        Node parent = n.parent;
        Node sibling = n.sibling();

        if(isRed(sibling)) {
            // case3 兄弟为红色
            if(n.isLeftChild()) {
                leftRotate(parent);
            } else {
                rightRotate(parent);
            }
            parent.color = RED;
            sibling.color = BLACK;

            fixDoubleBlack(n);
        } else {
            if(sibling != null) {
                if(isBlack(sibling.left) && isBlack(sibling.right)) {
                    // case4 兄弟是黑，两个侄子也是黑
                    sibling.color = RED;
                    if(isRed(parent)) {
                        parent.color = BLACK;
                    } else {
                        fixDoubleBlack(parent);
                    }
                } else {
                    // case5 兄弟为黑，两个侄子有一个或都是红
                    if(sibling.isLeftChild() && isRed(sibling.left)) { // LL
                        rightRotate(parent);
                        sibling.left.color = BLACK;
                        sibling.color = parent.color;
                    } else if(sibling.isLeftChild() && isRed(sibling.right)) { // LR
                        sibling.right.color = parent.color;
                        leftRotate(sibling);
                        rightRotate(parent);
                    } else if (!sibling.isLeftChild() && isRed(sibling.left)) { // RL
                        sibling.left.color = parent.color;
                        rightRotate(sibling);
                        leftRotate(parent);
                    } else { // RR
                        leftRotate(parent);
                        sibling.right.color = BLACK;
                        sibling.color = parent.color;
                    }
                    parent.color = BLACK;
                }
            } else {
                fixDoubleBlack(parent);
            }
        }
    }

    private Node findReplaced(Node n) {
        if(n.left == null && n.right == null) {
            return null;
        }
        if(n.left == null) {
            return n.right;
        }
        if(n.right == null) {
            return n.left;
        }

        // 找后继节点
        Node s = n.right;
        while (s.left != null) {
            s = s.left;
        }
        return s;
    }

    private Node find(int key) {
        Node p = root;
        while (p != null) {
            if (key < p.key) {
                p = p.left;
            } else if (key > p.key) {
                p = p.right;
            } else {
                return p;
            }
        }
        return null;
    }

    /**
     * 新增
     */
    public void put(int key, Object value) {
        Node p = root;
        Node parent = null;
        while (p != null) {
            parent = p;
            if(key < p.key) {
                p = p.left;
            } else if (key > p.key) {
                p = p.right;
            } else {
                // 更新逻辑
                p.value = value;
                return;
            }
        }

        Node inserted = new Node(key, value);
        if(parent == null) {
            root = inserted;
        } else if(key < parent.key) {
            parent.left = inserted;
            inserted.parent = parent;
        } else {
            parent.right = inserted;
            inserted.parent = parent;
        }

        fixRedRed(inserted);
    }

    /**
     * 新增的4种情况处理，默认插入节点 n 为红色，这里的n必不可能为黑色
     * case1：n为根节点 => 根节点变黑
     * case2：n的父节点为黑色 => 正常插入，不用其它处理
     * case3：n的父亲是红色，叔叔是红色 => 先把父亲和叔叔变黑，爷爷变红，爷爷递归调整（又会经过case1，case2判断）
     * case4：n的父亲是红色，叔叔是黑色 =>
     *      父亲是左孩子，n是左孩子 LL   父亲变黑，爷爷变红，爷爷右旋
     *      父亲是左孩子，n是右孩子 LR   父亲左旋，n变黑，爷爷变红，爷爷右旋
     *      父亲是右孩子，n是左孩子 RL   父亲右旋，n变黑，爷爷变红，爷爷左旋
     *      父亲是右孩子，n是右孩子 RR   父亲变黑，爷爷变红，爷爷左旋
     */
    private void fixRedRed(Node n) {
//        assert n.color == RED;
        if(root == n) {
            root.color = BLACK;
            return;
        }

        if(isBlack(n.parent)) {
            return;
        }

//        assert isRed(n.parent);
//        assert n.parent != null;

        Node parent = n.parent;
        Node uncle = n.uncle();
        Node grandparent = parent.parent;

        // 红红相邻处理
        if(isRed(uncle)) {
//            assert uncle != null; // isRed(uncle)返回true，那么uncle必不为null
            parent.color = BLACK;
            uncle.color = BLACK;
            // 红红的话，该节点最低在第3层，因为根节点是黑色，所以grandparent不可能为mull
            grandparent.color = RED;
            fixRedRed(grandparent);
        } else {
            if(parent.isLeftChild() && n.isLeftChild()) { // LL  父亲是左孩子，n也是左孩子
                parent.color = BLACK;
                grandparent.color = RED;
                rightRotate(grandparent);
            } else if(parent.isLeftChild()) { // LR
                leftRotate(parent);
                n.color = BLACK;
                grandparent.color = RED;
                rightRotate(grandparent);
            } else if(n.isLeftChild()) { // RL
                rightRotate(parent);
                n.color = BLACK;
                grandparent.color = RED;
                leftRotate(grandparent);
            } else { // RR
                parent.color = BLACK;
                grandparent.color = RED;
                leftRotate(grandparent);
            }
        }
    }

    // 左旋
    void leftRotate(Node pink) {
        Node parent = pink.parent;

        Node yellow = pink.right;
        Node green = yellow.left;
        yellow.left = pink;
        pink.right = green;

        yellow.parent = parent;
        pink.parent = yellow;
        if(green != null) {
            green.parent = pink;
        }

        if(parent == null) {
            root = yellow;
        } else if(parent.left == pink) {
            parent.left = yellow;
        } else {
            parent.right = yellow;
        }
    }

    // 右旋
    void rightRotate(Node pink) {
        Node parent = pink.parent;

        Node yellow = pink.left;
        Node green = yellow.right;
        yellow.right = pink;
        pink.left = green;

        // 换爹
        yellow.parent = parent;
        pink.parent = yellow;
        if (green != null) {
            green.parent = pink;
        }

        if(parent == null) {
            root = yellow;
        } else if(parent.left == pink) {
            parent.left = yellow;
        } else {
            parent.right = yellow;
        }
    }

    // 判红，非红即黑
    private boolean isRed(Node n) {
        return n != null && n.color == RED;
    }

    private boolean isBlack(Node n) {
        return n == null || n.color == BLACK;
    }

    public boolean contains(int k) {
        return find(k) != null;
    }

    public Object get(int k) {
        Node n = find(k);
        return n == null ? null: n.value;
    }


    public void printTree() {
        StringBuilder sb = new StringBuilder();
        printVerticalTree(root, "", sb);
        System.out.println(sb);
    }

    private void printVerticalTree(Node node, String prefix, StringBuilder sb) {
        if (node == null) {
            return;
        }

        printVerticalTree(node.right, prefix + "        ", sb);
        sb.append(node.color == RED ? "\u001B[31m" : "\u001B[30m");
        sb.append(prefix).append("|--");
        sb.append(node.key).append("\u001B[0m").append("\n");
        printVerticalTree(node.left, prefix + "        ", sb);
    }

    long redCount;
    long blackCount;

    /**
     * 打印红黑节点个数
     */
    public void printRBCount() {
        printCount(root);
        System.out.println("插入完毕，红色个数：" + redCount + ",黑色个数：" + blackCount);
    }

    private void printCount(Node node) {
        if (node == null) {
            return;
        }
        if(node.color == RED) {
            redCount++;
        } else {
            blackCount++;
        }
        printCount(node.right);
        printCount(node.left);
    }

}
